{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The PyNE `mesh` module is a part of the PyNE library, which provides functionalities for creating and manipulating generic meshes. It supports both cartesian and tetrahedral meshes. The meshes created using PyNE can be tagged with scalar and vector quantities, as well as PyNE `Material` objects. These tagged data can be easily accessed using a common syntax.\n",
    "\n",
    "Internally, the PyNE `Mesh` class is a wrapper around a mesh library called [MOAB (Mesh Oriented datABase)](https://sigma.mcs.anl.gov/moab-library/). A `Mesh` object in PyNE stores a mesh instance from MOAB and also includes a PyNE `MaterialLibrary` object, which is used to map PyNE `Material`s to the volume elements of the mesh.\n",
    "\n",
    "In this notebook, we will explore some basic examples of using PyNE meshes. We will begin by importing the necessary components from PyNE and other related libraries, which will be explained further in the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Known error: This notebook requires [yt](https://yt-project.org/) to visualize the results. Yt needs to be updated to work properly first. Updates are currently being made to yt's frontends to make this PyNE integration work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import some classes from pyne.mesh\n",
    "from pyne.mesh import Mesh, NativeMeshTag, MetadataTag, ComputedTag\n",
    "from pyne.material import Material, from_atom_frac\n",
    "from pyne.xs.channels import sigma_t\n",
    "import numpy as np\n",
    "import logging\n",
    "\n",
    "# Disable all logging outputs\n",
    "logging.disable(logging.CRITICAL)\n",
    "\n",
    "# import plotting utilities\n",
    "from yt.frontends.moab.api import PyneMoabHex8Dataset\n",
    "from yt import SlicePlot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make a simple Cartesian mesh that is 4 x 3 x 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "coords = [[0, 1, 2, 3, 4], [0, 1, 2, 3], [0,1]]\n",
    "m = Mesh(structured=True, structured_coords=coords) # here \"structured\" means Cartesian"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tags, also referred to as fields, are utilized to store data on a mesh. PyNE `NativeMeshTags` enables the tagging of both scalar and vector quantities on each volume element of the mesh. Let's consider an example where we create a tag to represent the neutron flux on a mesh. This tag is a floating-point scalar value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.neutron_flux = NativeMeshTag(size=1, dtype=float)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The mesh can now be populated with fictitious neutron flux data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.neutron_flux[:] = [0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 18.0, 20.0, 22.0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyNE meshes are supported by the [yt project](http://yt-project.org/).  Here we use yt's plotting infrastructure to display the mesh tag that was just created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img style=\"max-width:100%;max-height:100%;\" src=\"\"><br>"
      ],
      "text/plain": [
       "<yt.visualization.plot_window.AxisAlignedSlicePlot at 0x7f68b8be1390>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pf = PyneMoabHex8Dataset(m)\n",
    "s = SlicePlot(pf, 'z', 'neutron_flux', origin='native')\n",
    "s.set_cmap('neutron_flux', 'Purples')\n",
    "s.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The volume elements within the mesh are kept track of using the volume element index (`i`). The volume element index is defined on the range from 0 (inclusive) to the number of volumes in the mesh (exclusive).\n",
    "\n",
    "The flux for the volume element `i=7` can be retrieved (and similarly set) using the following syntax:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([14.])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.neutron_flux[7]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"Fancy\" indexing can be used. This means that multiple indicies can be provided, yielding the tag value for each:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([4., 6., 8.])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.neutron_flux[2, 3, 4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Masking can be used to return only values that match some criteria. Here, we request all flux values above 10.0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([10., 12., 14., 16., 18., 20., 22.])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.neutron_flux[m.neutron_flux[:] >= 10.0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyNE Mesh + PyNE Materials"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In PyNE, meshes have the capability to natively store PyNE `Material` objects, which allows for the representation of discretized versions of geometries. To illustrate this, let's create a new mesh called `m2`, which is slightly larger than our previous example. Specifically, `m2` has dimensions of 9x9x1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "coords2 = [np.linspace(0,9,10), np.linspace(0,9,10), [0, 10]]\n",
    "m2 = Mesh(structured=True, structured_coords=coords2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a few materials that will be assigned to the mesh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "fuel = from_atom_frac({'U235': 0.045, 'U238': 0.955, 'O16': 2.0}, density=10.7)\n",
    "mod = from_atom_frac({'H1': 2.0, 'O16': 1.0}, density=1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loop through the mesh and assign materials in a reactor-like pattern."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(m2)):\n",
    "    if (i+1)%2 != 0 and np.floor((i)/9)%2 != 0:\n",
    "        m2.mats[i] = fuel\n",
    "    else:    \n",
    "        m2.mats[i] = mod"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the density. Here, the fuel appears red (high density) and the moderator appears blue (low density)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img style=\"max-width:100%;max-height:100%;\" src=\"\"><br>"
      ],
      "text/plain": [
       "<yt.visualization.plot_window.AxisAlignedSlicePlot at 0x7f68b5505710>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pf2 = PyneMoabHex8Dataset(m2)\n",
    "s2 = SlicePlot(pf2, 'z', 'density', origin='native')\n",
    "s2.set_cmap('density', 'RdBu_r')\n",
    "s2.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The materials that we've added to the mesh can be accesssed via the same syntax used for `NativeMeshTags`. The material on the the mesh volume element with index `i=12` is retrieved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Material:\n",
      "mass = 269.90530758893\n",
      "density = 10.7\n",
      "atoms per molecule = 3.0\n",
      "mat_number = 13\n",
      "name = 12\n",
      "------------------------\n",
      "O16    0.11852241634284943\n",
      "U235   0.03918773166424315\n",
      "U238   0.8422898519929074\n"
     ]
    }
   ],
   "source": [
    "print(m2.mats[12])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The attributes of this material can be accessed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10.7"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m2.mats[12].density"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Equivalently, the density can be accessed via the generic tag syntax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m2.density[3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same is true of the composition of the material (another attribute of the `Material` class)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{10010000: 0.11191487328888054, 80160000: 0.8880851267111195}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m2.comp[3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In fact, all `Material` attributes are treated as tags. A list of all tags, including these material properties treated as tags, is shown here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NativeMeshTag: idx\n",
      "MaterialPropertyTag: atoms_per_molecule\n",
      "MaterialPropertyTag: metadata\n",
      "MaterialPropertyTag: comp\n",
      "MaterialPropertyTag: mass\n",
      "MaterialPropertyTag: density\n",
      "MaterialMethodTag: expand_elements\n",
      "MaterialMethodTag: mass_density\n",
      "MaterialMethodTag: molecular_mass\n",
      "MaterialMethodTag: mult_by_mass\n",
      "MaterialMethodTag: number_density\n",
      "MaterialMethodTag: sub_act\n",
      "MaterialMethodTag: sub_fp\n",
      "MaterialMethodTag: sub_lan\n",
      "MaterialMethodTag: sub_ma\n",
      "MaterialMethodTag: sub_tru\n",
      "MaterialMethodTag: to_atom_frac\n"
     ]
    }
   ],
   "source": [
    "for key, value in m2.tags.items():\n",
    "    print(f\"{value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Advanced Feature: Computed Tags"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A computed tag, also known as a derived field, is a method of creating a *virtual* tag within a system. Unlike regular tags, computed tags do not require the storage of tag data directly on the mesh or pre-computed values. Instead, computed tags calculate and return values on-the-fly when requested, using a technique known as lazy evaluation. This approach is particularly useful when dealing with large datasets.\n",
    "\n",
    "To define a computed tag, you need to provide a function, lambda (also called an anonymous function), or any other callable object that meets the following interface requirements:\n",
    "\n",
    "```python\n",
    "def f(mesh, i):\n",
    "    \"\"\"mesh is a pyne.mesh.Mesh() object and i is the volume element index\n",
    "    to compute.\n",
    "    \"\"\"\n",
    "    # ... do some work ...\n",
    "    return anything_you_want\n",
    "```\n",
    "\n",
    "Let's consider a playful example that involves the creation of a computed tag called `i_plus_density`. In this scenario, `i_plus_density` is derived by summing the index `i` of a mesh volume element with its corresponding density value. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "m2.i_plus_density = ComputedTag(lambda mesh, i: i + mesh.density[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting `i_plus_density`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img style=\"max-width:100%;max-height:100%;\" src=\"\"><br>"
      ],
      "text/plain": [
       "<yt.visualization.plot_window.AxisAlignedSlicePlot at 0x7f68b552d7d0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pf3 = PyneMoabHex8Dataset(m2)\n",
    "s3 = SlicePlot(pf3, 'z', 'i_plus_density', origin='native')\n",
    "s3.set_cmap('i_plus_density', 'RdBu_r')\n",
    "s3.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a more serious example which uses PyNE's cross section tools to compute the one-group total cross section $\\sigma_t$ everywhere on the mesh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "m2.sigma_t = ComputedTag(lambda mesh, i: sigma_t(mesh.mats[i], group_struct=[10.0, 1e-6], phi_g=[1.0])[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img style=\"max-width:100%;max-height:100%;\" src=\"\"><br>"
      ],
      "text/plain": [
       "<yt.visualization.plot_window.AxisAlignedSlicePlot at 0x7f68b4ed3610>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pf4 = PyneMoabHex8Dataset(m2)\n",
    "s4 = SlicePlot(pf4, 'z', 'sigma_t', origin='native')\n",
    "s4.set_cmap('sigma_t', 'RdBu_r')\n",
    "s4.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This plot demonstrates that the one-group total cross section reaches its peak within the moderator. Remarkably, this insightful observation was generated using just a single line of code!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
